package com.agent.server.service;

import com.agent.server.sse.LlmSseClient;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Service;
import reactor.core.Disposable;

import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.function.Consumer;

@Service
public class LlmService {
    private static final Logger logger = LoggerFactory.getLogger(LlmService.class);
    
    @Value("${llm.websocket.uri:http://localhost:8081}")
    private String llmBaseUrl;
    
    @Value("${llm.websocket.secret-token:}")
    private String secretToken;
    
    private final ConcurrentHashMap<String, Disposable> activeSubscriptions = new ConcurrentHashMap<>();
    
    /**
     * Create a new connection to the LLM service and stream results back to the gRPC client
     *
     * @param sessionId Unique identifier for this conversation session
     * @param prompt The user's prompt/query
     * @param messageConsumer Consumer that will handle the LLM responses
     * @return true if connection was successfully established, false otherwise
     */
    public boolean streamFromLlm(String sessionId, String prompt, Consumer<String> messageConsumer) {
        try {
            // Close any existing connection for this session
            closeConnection(sessionId);
            
            // Set up headers for authentication
            Map<String, String> headers = new HashMap<>();
            if (!secretToken.isEmpty()) {
                headers.put("Authorization", "Bearer " + secretToken);
            }
            
            // Create a new SSE client
            LlmSseClient client = new LlmSseClient(llmBaseUrl, headers);
            
            // Subscribe to the SSE stream
            Disposable subscription = client.streamRequestWithConsumer(prompt, messageConsumer)
                .subscribe(
                    response -> {
                        // Response processing is handled by the consumer
                        logger.debug("Response processed by consumer");
                    },
                    error -> {
                        logger.error("Error in SSE stream", error);
                        activeSubscriptions.remove(sessionId);
                    },
                    () -> {
                        logger.info("SSE stream completed for session: {}", sessionId);
                        activeSubscriptions.remove(sessionId);
                    }
                );
            
            // Store the subscription
            activeSubscriptions.put(sessionId, subscription);
            logger.info("Started LLM streaming for session: {}", sessionId);
            return true;
        } catch (Exception e) {
            logger.error("Failed to create SSE stream", e);
            return false;
        }
    }
    
    /**
     * Close an active SSE connection
     * 
     * @param sessionId The session ID to close
     */
    public void closeConnection(String sessionId) {
        Disposable subscription = activeSubscriptions.remove(sessionId);
        if (subscription != null && !subscription.isDisposed()) {
            subscription.dispose();
            logger.info("Closed SSE connection for session: {}", sessionId);
        }
    }
    
    /**
     * Close all active SSE connections
     */
    public void closeAllConnections() {
        activeSubscriptions.forEach((sessionId, subscription) -> {
            if (!subscription.isDisposed()) {
                subscription.dispose();
                logger.info("Closed SSE connection for session: {}", sessionId);
            }
        });
        activeSubscriptions.clear();
    }
} 